package cn.genmer.test.security.machinelearning.deeplearning4j.text;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.Updater;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.EmbeddingLayer;
import org.deeplearning4j.nn.conf.layers.GravesLSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;

/**
 *  训练数据 https://spaces.ac.cn/archives/3414
 */
public class TextTrain {
    public static final String BASE_PATH = "/Users/genmer/Documents/Codes/tensorFlowModel/mnist";
    public static void train() throws Exception {

        int batchSize = 32;
        int numEpochs = 5;
        int seed = 12345;
        // 设置训练数据和测试数据迭代器
        DataSetIterator trainIter = new MnistDataSetIterator(batchSize, true, seed);
        DataSetIterator testIter = new MnistDataSetIterator(batchSize, false, seed);


        // 构建神经网络的配置
        MultiLayerNetwork model = textClassifyModel();
        System.out.println("模型配置信息：" + model.getLayerWiseConfigurations());
        model.init();

        model.setListeners(new ScoreIterationListener(1));
        // 训练模型
        for (int i = 0; i < numEpochs; i++) {
            model.fit(trainIter);
        }

        // 在测试数据上评估模型
        Evaluation eval = model.evaluate(testIter);
        System.out.println(eval.stats());

        // 将模型保存到本地磁盘
        File locationToSave = new File(BASE_PATH+"/cnn_mnist_model.zip");
        boolean saveUpdater = true; // 是否保存updater（用于进行模型参数更新）
        ModelSerializer.writeModel(model, locationToSave, saveUpdater);
    }

    public static MultiLayerNetwork textClassifyModel(){
        Integer VOCAB_SIZE = 0;
        MultiLayerConfiguration netconf = new NeuralNetConfiguration.Builder()
                .seed(1234)
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .l2(5 * 1e-4)
                .updater(Updater.ADAM)
                .list()
                .layer(0, new EmbeddingLayer.Builder().nIn(VOCAB_SIZE).nOut(100).activation(Activation.IDENTITY).build())
                .layer(1, new GravesLSTM.Builder().nIn(100).nOut(100).activation(Activation.SOFTSIGN).build())
                .layer(2, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MCXENT)
                        .activation(Activation.SOFTMAX).nIn(100).nOut(2).build())
                .setInputType(InputType.recurrent(VOCAB_SIZE))
                .build();

        MultiLayerNetwork net = new MultiLayerNetwork(netconf);
        return net;
    }

    public static void main(String[] args) throws Exception {
        train();
    }
}